{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "0",
   "metadata": {},
   "source": [
    "# 8. OpenAI Responses Target\n",
    "\n",
    "In this demo, we show an example of the `OpenAIResponseTarget`. [Responses](https://platform.openai.com/docs/api-reference/responses) is a newer protocol than chat completions and provides additional functionality with a somewhat modified API. The allowed input types include text, image, web search, file search, functions, reasoning, and computer use.\n",
    "\n",
    "Before you begin, ensure you are set up with the correct version of PyRIT installed and have secrets configured as described [here](../../setup/populating_secrets.md).\n",
    "\n",
    "\n",
    "## OpenAI Configuration\n",
    "\n",
    "Like most targets, all `OpenAITarget`s need an `endpoint` and often also needs a `model` and a `key`. These can be passed into the constructor or configured with environment variables (or in .env).\n",
    "\n",
    "- endpoint: The API endpoint (`OPENAI_RESPONSES_ENDPOINT` environment variable). For OpenAI, these are just \"https://api.openai.com/v1/responses\".\n",
    "- auth: The API key for authentication (`OPENAI_RESPONSES_KEY` environment variable).\n",
    "- model_name: The model to use (`OPENAI_RESPONSES_MODEL` environment variable). For OpenAI, these are any available model name and are listed here: \"https://platform.openai.com/docs/models\"."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "\u001b[34m────────────────────────────────────────────────────────────────────────────────────────────────────\u001b[0m\n",
      "\u001b[1m\u001b[34m🔹 Turn 1 - USER\u001b[0m\n",
      "\u001b[34m────────────────────────────────────────────────────────────────────────────────────────────────────\u001b[0m\n",
      "\u001b[34m  Tell me a joke\u001b[0m\n",
      "\n",
      "\u001b[33m────────────────────────────────────────────────────────────────────────────────────────────────────\u001b[0m\n",
      "\u001b[1m\u001b[33m🔸 ASSISTANT\u001b[0m\n",
      "\u001b[33m────────────────────────────────────────────────────────────────────────────────────────────────────\u001b[0m\n",
      "\u001b[33m  Why don’t skeletons fight each other?\u001b[0m\n",
      "\u001b[33m    They don’t have the guts.\u001b[0m\n",
      "\n",
      "\u001b[34m────────────────────────────────────────────────────────────────────────────────────────────────────\u001b[0m\n"
     ]
    }
   ],
   "source": [
    "from pyrit.executor.attack import ConsoleAttackResultPrinter, PromptSendingAttack\n",
    "from pyrit.prompt_target import OpenAIResponseTarget\n",
    "from pyrit.setup import IN_MEMORY, initialize_pyrit\n",
    "\n",
    "initialize_pyrit(memory_db_type=IN_MEMORY)\n",
    "\n",
    "target = OpenAIResponseTarget()\n",
    "# For an AzureOpenAI endpoint with Entra ID authentication enabled, use the following command instead. Make sure to run `az login` first.\n",
    "# target = OpenAIResponseTarget(use_entra_auth=True)\n",
    "\n",
    "attack = PromptSendingAttack(objective_target=target)\n",
    "\n",
    "result = await attack.execute_async(objective=\"Tell me a joke\")  # type: ignore\n",
    "await ConsoleAttackResultPrinter().print_conversation_async(result=result)  # type: ignore"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2",
   "metadata": {},
   "source": [
    "## Tool Use with Custom Functions\n",
    "\n",
    "In this example, we demonstrate how the OpenAI `Responses API` can be used to invoke a **custom-defined Python function** during a conversation. This is part of OpenAI’s support for \"function calling\", where the model decides to call a registered function, and the application executes it and passes the result back into the conversation loop.\n",
    "\n",
    "We define a simple tool called `get_current_weather`, which simulates weather information retrieval. A corresponding OpenAI tool schema describes the function name, parameters, and expected input format.\n",
    "\n",
    "The function is registered in the `custom_functions` argument of `OpenAIResponseTarget`. The `extra_body_parameters` include:\n",
    "\n",
    "- `tools`: the full OpenAI tool schema for `get_current_weather`.\n",
    "- `tool_choice: \"auto\"`: instructs the model to decide when to call the function.\n",
    "\n",
    "The user prompt explicitly asks the model to use the `get_current_weather` function. Once the model responds with a `function_call`, PyRIT executes the function, wraps the output, and the conversation continues until a final answer is produced.\n",
    "\n",
    "This showcases how agentic function execution works with PyRIT + OpenAI Responses API."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "0 | assistant: {\"id\":\"rs_00b282377b838d1b0069320419a68881958724cfda02484cfb\",\"type\":\"reasoning\",\"summary\":[],\"content\":null,\"encrypted_content\":null}\n",
      "1 | assistant: {\"type\":\"function_call\",\"call_id\":\"call_d068M8MPsIKWcS1z7ckOHhR2\",\"name\":\"get_current_weather\",\"arguments\":\"{\\\"location\\\":\\\"Boston\\\",\\\"unit\\\":\\\"celsius\\\"}\"}\n",
      "0 | tool: {\"type\":\"function_call_output\",\"call_id\":\"call_d068M8MPsIKWcS1z7ckOHhR2\",\"output\":\"{\\\"weather\\\":\\\"Sunny\\\",\\\"temp_c\\\":22,\\\"location\\\":\\\"Boston\\\",\\\"unit\\\":\\\"celsius\\\"}\"}\n",
      "0 | assistant: The current weather in Boston is Sunny with a temperature of 22°C.\n"
     ]
    }
   ],
   "source": [
    "from pyrit.models import Message, MessagePiece\n",
    "from pyrit.setup import IN_MEMORY, initialize_pyrit\n",
    "\n",
    "initialize_pyrit(memory_db_type=IN_MEMORY)\n",
    "\n",
    "\n",
    "async def get_current_weather(args):\n",
    "    return {\n",
    "        \"weather\": \"Sunny\",\n",
    "        \"temp_c\": 22,\n",
    "        \"location\": args[\"location\"],\n",
    "        \"unit\": args[\"unit\"],\n",
    "    }\n",
    "\n",
    "\n",
    "# Responses API function tool schema (flat, no nested \"function\" key)\n",
    "function_tool = {\n",
    "    \"type\": \"function\",\n",
    "    \"name\": \"get_current_weather\",\n",
    "    \"description\": \"Get the current weather in a given location\",\n",
    "    \"parameters\": {\n",
    "        \"type\": \"object\",\n",
    "        \"properties\": {\n",
    "            \"location\": {\"type\": \"string\", \"description\": \"City and state\"},\n",
    "            \"unit\": {\"type\": \"string\", \"enum\": [\"celsius\", \"fahrenheit\"]},\n",
    "        },\n",
    "        \"required\": [\"location\", \"unit\"],\n",
    "        \"additionalProperties\": False,\n",
    "    },\n",
    "    \"strict\": True,\n",
    "}\n",
    "\n",
    "# Let the model auto-select tools\n",
    "target = OpenAIResponseTarget(\n",
    "    custom_functions={\"get_current_weather\": get_current_weather},\n",
    "    extra_body_parameters={\n",
    "        \"tools\": [function_tool],\n",
    "        \"tool_choice\": \"auto\",\n",
    "    },\n",
    "    httpx_client_kwargs={\"timeout\": 60.0},\n",
    ")\n",
    "\n",
    "# Build the user prompt\n",
    "message_piece = MessagePiece(\n",
    "    role=\"user\",\n",
    "    original_value=\"What is the weather in Boston in celsius? Use the get_current_weather function.\",\n",
    "    original_value_data_type=\"text\",\n",
    ")\n",
    "message = Message(message_pieces=[message_piece])\n",
    "\n",
    "response = await target.send_prompt_async(message=message)  # type: ignore\n",
    "\n",
    "for response_msg in response:\n",
    "    for idx, piece in enumerate(response_msg.message_pieces):\n",
    "        print(f\"{idx} | {piece.role}: {piece.original_value}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4",
   "metadata": {},
   "source": [
    "## Using the Built-in Web Search Tool\n",
    "\n",
    "In this example, we use a built-in PyRIT helper function `web_search_tool()` to register a web search tool with OpenAI's Responses API. This allows the model to issue web search queries during a conversation to supplement its responses with fresh information.\n",
    "\n",
    "The tool is added to the `extra_body_parameters` passed into the `OpenAIResponseTarget`. As before, `tool_choice=\"auto\"` enables the model to decide when to invoke the tool.\n",
    "\n",
    "The user prompt asks for a recent positive news story — an open-ended question that may prompt the model to issue a web search tool call. PyRIT will automatically execute the tool and return the output to the model as part of the response.\n",
    "\n",
    "This example demonstrates how retrieval-augmented generation (RAG) can be enabled in PyRIT through OpenAI's Responses API and integrated tool schema.\n",
    "\n",
    "NOTE that web search is NOT supported through an Azure OpenAI endpoint, only through the OpenAI platform endpoint (i.e. api.openai.com)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "0 | assistant: {\"type\":\"web_search_call\",\"id\":\"ws_0d9a4535c6aeb49e00693204204e488197a7e84cc5b152880e\"}\n",
      "1 | assistant: One positive news story from today is about a remarkable breakthrough in healthcare: Moderna’s mRNA cancer vaccine, now in Phase 3 trials, has shown a 44% reduction in melanoma recurrence, bringing new hope to millions affected by this disease. This advance is being celebrated as a potential game-changer in cancer treatment in 2025 [Positive News That Mattered: Breakthroughs in Science, Tech, and Humanity in 2025](https://www.iuemag.com/g25/sc/positive-news-that-mattered-breakthroughs-in-science-tech-and-humanity-in-2025).\n"
     ]
    }
   ],
   "source": [
    "import os\n",
    "\n",
    "from pyrit.common.tool_configs import web_search_tool\n",
    "from pyrit.models import Message, MessagePiece\n",
    "from pyrit.prompt_target import OpenAIResponseTarget\n",
    "from pyrit.setup import IN_MEMORY, initialize_pyrit\n",
    "\n",
    "initialize_pyrit(memory_db_type=IN_MEMORY)\n",
    "\n",
    "# Note: web search is only supported on a limited set of models.\n",
    "target = OpenAIResponseTarget(\n",
    "    endpoint=os.getenv(\"AZURE_OPENAI_GPT41_RESPONSES_ENDPOINT\"),\n",
    "    api_key=os.getenv(\"AZURE_OPENAI_GPT41_RESPONSES_KEY\"),\n",
    "    model_name=os.getenv(\"AZURE_OPENAI_GPT41_RESPONSES_MODEL\"),\n",
    "    extra_body_parameters={\n",
    "        \"tools\": [web_search_tool()],\n",
    "        \"tool_choice\": \"auto\",\n",
    "    },\n",
    "    httpx_client_kwargs={\"timeout\": 60},\n",
    ")\n",
    "\n",
    "message_piece = MessagePiece(\n",
    "    role=\"user\", original_value=\"Briefly, what is one positive news story from today?\", original_value_data_type=\"text\"\n",
    ")\n",
    "message = Message(message_pieces=[message_piece])\n",
    "\n",
    "response = await target.send_prompt_async(message=message)  # type: ignore\n",
    "\n",
    "for response_msg in response:\n",
    "    for idx, piece in enumerate(response_msg.message_pieces):\n",
    "        print(f\"{idx} | {piece.role}: {piece.original_value}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6",
   "metadata": {},
   "source": [
    "## Grammar-Constrained Generation\n",
    "\n",
    "OpenAI models also support constrained generation in the [Responses API](https://platform.openai.com/docs/guides/function-calling#context-free-grammars). This forces the LLM to produce output which conforms to the given grammar, which is useful when specific syntax is required in the output.\n",
    "\n",
    "In this example, we will define a simple Lark grammar which prevents the model from giving a correct answer to a simple question, and compare that to the unconstrained model.\n",
    "\n",
    "Note that as of October 2025, this is only supported by OpenAI (not Azure) on \"gpt-5\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Unconstrained Response:\n",
      "0 | assistant: {\"id\":\"rs_0d7e8ef21bbbd586006932042513e08195b16944b6419aa1f5\",\"type\":\"reasoning\",\"summary\":[],\"content\":null,\"encrypted_content\":null}\n",
      "1 | assistant: Rome.\n",
      "\n",
      "Constrained Response:\n",
      "0 | assistant: {\"id\":\"rs_0a3146a3baf9205a006932042b08748190b3d015fd75ab1803\",\"type\":\"reasoning\",\"summary\":[],\"content\":null,\"encrypted_content\":null}\n",
      "1 | assistant: I think that it is Padua\n"
     ]
    }
   ],
   "source": [
    "from pyrit.setup import IN_MEMORY, initialize_pyrit\n",
    "\n",
    "initialize_pyrit(memory_db_type=IN_MEMORY)\n",
    "\n",
    "message_piece = MessagePiece(\n",
    "    role=\"user\",\n",
    "    original_value=\"What is the capital of Italy?\",\n",
    "    original_value_data_type=\"text\",\n",
    ")\n",
    "message = Message(message_pieces=[message_piece])\n",
    "\n",
    "# Define a grammar that prevents \"Rome\" from being generated\n",
    "lark_grammar = r\"\"\"\n",
    "start: \"I think that it is \" SHORTTEXT\n",
    "SHORTTEXT: /[^RrOoMmEe]{1,8}/\n",
    "\"\"\"\n",
    "\n",
    "grammar_tool = {\n",
    "    \"type\": \"custom\",\n",
    "    \"name\": \"CitiesGrammar\",\n",
    "    \"description\": \"Constrains generation.\",\n",
    "    \"format\": {\n",
    "        \"type\": \"grammar\",\n",
    "        \"syntax\": \"lark\",\n",
    "        \"definition\": lark_grammar,\n",
    "    },\n",
    "}\n",
    "\n",
    "target = OpenAIResponseTarget(\n",
    "    endpoint=os.getenv(\"AZURE_OPENAI_GPT5_RESPONSES_ENDPOINT\"),\n",
    "    api_key=os.getenv(\"AZURE_OPENAI_GPT5_KEY\"),\n",
    "    model_name=os.getenv(\"AZURE_OPENAI_GPT5_MODEL\"),\n",
    "    extra_body_parameters={\"tools\": [grammar_tool], \"tool_choice\": \"required\"},\n",
    "    temperature=1.0,\n",
    ")\n",
    "\n",
    "unconstrained_target = OpenAIResponseTarget(\n",
    "    endpoint=os.getenv(\"AZURE_OPENAI_GPT5_RESPONSES_ENDPOINT\"),\n",
    "    api_key=os.getenv(\"AZURE_OPENAI_GPT5_KEY\"),\n",
    "    model_name=os.getenv(\"AZURE_OPENAI_GPT5_MODEL\"),\n",
    "    temperature=1.0,\n",
    ")\n",
    "\n",
    "unconstrained_result = await unconstrained_target.send_prompt_async(message=message)  # type: ignore\n",
    "\n",
    "result = await target.send_prompt_async(message=message)  # type: ignore\n",
    "\n",
    "print(\"Unconstrained Response:\")\n",
    "for response_msg in unconstrained_result:\n",
    "    for idx, piece in enumerate(response_msg.message_pieces):\n",
    "        print(f\"{idx} | {piece.role}: {piece.original_value}\")\n",
    "\n",
    "print()\n",
    "\n",
    "print(\"Constrained Response:\")\n",
    "for response_msg in result:\n",
    "    for idx, piece in enumerate(response_msg.message_pieces):\n",
    "        print(f\"{idx} | {piece.role}: {piece.original_value}\")"
   ]
  }
 ],
 "metadata": {
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
